package com.yuxl.ai.knowledgeText.service;

import com.alibaba.fastjson2.JSONArray;
import com.yuxl.ai.chat.domain.ReRankReq;
import com.yuxl.ai.chat.domain.ReRankRes;
import com.yuxl.ai.chat.service.ReRankApi;
import com.yuxl.ai.embedding.service.EmbeddingApi;
import com.yuxl.ai.knowledgeText.domain.AiKnowledgeTextEmbeddingQuery;
import com.yuxl.ai.knowledgeText.domain.AiKnowledgeTextEmbeddingRes;
import com.yuxl.ai.knowledgeText.entity.AiKnowledgeTextEmbedding;
import com.yuxl.ai.knowledgeText.repository.AiKnowledgeTextEmbeddingRepository;
import com.yuxl.ai.knowledgeText.repository.AiKnowledgeTextRepository;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * @Author : yuxl
 * @create 2025/7/11 上午9:43
 */
@Slf4j
@Service
public class AiKnowledgeTextEmbeddingService {
    @Resource
    private AiKnowledgeTextEmbeddingRepository aiKnowledgeTextEmbeddingRepository;
    @Resource
    private EmbeddingApi embeddingApi;
    @Resource
    private ReRankApi reRankApi;
    @Resource
    private AiKnowledgeTextRepository aiKnowledgeTextRepository;

    public static float[] l2Normalize(float[] vector) {
        double norm = Math.sqrt(IntStream.range(0, vector.length).mapToDouble(i -> vector[i] * vector[i]).sum());
        if (norm == 0) {
            return vector;
        }
        float[] normalized = new float[vector.length];
        for (int i = 0; i < vector.length; i++) {
            normalized[i] = (float) (vector[i] / norm);
        }
        return normalized;
    }

    public Flux<AiKnowledgeTextEmbedding> embedding(Long id) {
        return aiKnowledgeTextRepository.findById(id)
                .flatMap(entity -> {
                    List<String> contents = JSONArray.parseArray(entity.getFormat(), String.class);
                    return embeddingApi.embedding(contents)
                            .map(embeddings ->
                                    IntStream.range(0, embeddings.size())
                                            .mapToObj(idx -> new AiKnowledgeTextEmbedding()
                                                    .setTextId(id)
                                                    .setKnowledgeId(entity.getKnowledgeId())
                                                    .setName(entity.getName())
                                                    .setContent(contents.get(idx))
                                                    .setEmbedding(embeddings.get(idx))
                                                    .setCreateTime(LocalDateTime.now())
                                            ).toList()
                            );
                })
                .flatMapMany(list -> aiKnowledgeTextEmbeddingRepository.saveAll(list));
    }

    public Mono<List<AiKnowledgeTextEmbeddingRes>> query(AiKnowledgeTextEmbeddingQuery query) {
        int limit = query.getLimit();
        double similarity = query.getSimilarity();
        String content = query.getContent();
        return embeddingApi.embedding(query.getContent())
                .flatMapMany(embedding -> {
                    query.setContent(null);
                    query.setLimit(limit * 10);
                    query.setSimilarity(similarity / 10);
                    String embeddingStr = Arrays.toString(embedding);
                    return aiKnowledgeTextEmbeddingRepository.query(query, embeddingStr);
                })
                .collectList()
                .flatMap(list -> {
                    List<String> docs = list.stream().map(AiKnowledgeTextEmbeddingRes::getContent).toList();
                    ReRankReq q = new ReRankReq()
                            .setQuery(content)
                            .setTopK(limit)
                            .setTexts(docs);
                    return reRankApi.sort(q)
                            .map(outs -> {
                                Map<String, Double> map = outs.stream().filter(item -> item.getScore() > similarity).collect(Collectors.toMap(ReRankRes::getText, ReRankRes::getScore));
                                return list.stream()
                                        .filter(item -> map.containsKey(item.getContent()))
                                        .peek(item -> item.setSimilarity(map.get(item.getContent())))
                                        .sorted(Comparator.comparing(AiKnowledgeTextEmbeddingRes::getSimilarity).reversed())
                                        .limit(query.getLimit())
                                        .toList();
                            });
                });
    }
}
